#!/usr/bin/env python3

# Copyright (c) 2005-2014 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import glob
import stat
import pathlib
import tarfile
import subprocess
from ordered_set import OrderedSet
from ..._util import Util
from ..._util import TempChdir
from ..._po import HostMountPoint
from ..._po import HostDiskBtrfsRaid
from ..._po import HostDiskBcachefsRaid
from ..._po import HostDiskLvmLv
from ..._po import HostDiskBcache
from ..._po import HostDiskNvmeHdd
from ..._po import HostDiskScsiHdd
from ..._po import HostDiskXenHdd
from ..._po import HostDiskVirtioHdd
from ..._po import HostDiskImageFileOnRemovable
from ..._po import HostDiskFeatureUsePartitionTable
from ..._plugin import InitramfsPlugin


class Plugin(InitramfsPlugin):

    def modify_kernel_config(self, context):
        context.select_by_symbol("BLK_DEV_INITRD", "y")
        for sym_name in context.list_symbols_in_tree_by_symbol("BLK_DEV_INITRD", include_root=False):
            if sym_name in ["RD_XZ", "RD_LZMA"]:               # it seems RD_XZ has no effect, we have to also enable RD_LZMA, kernel bug?
                context.select_by_symbol(sym_name, "y")
            elif sym_name == "INITRAMFS_SOURCE":
                context.set_string_by_symbol(sym_name, "")
            else:
                context.select_by_symbol(sym_name, "n", allow_invisible=True)

    def install(self, context):
        self._be = context.get_boot_entry()
        self._initramfsTmpDir = "."                             # FIXME: redundant

        with TempChdir(context.get_tmp_dir()):
            subprocess.check_call("cp %s/* ." % (self.plugin_files_dir), shell=True)
            subprocess.check_call("make %s" % (context.get_build_variable("MAKEOPTS")), shell=True)

        if not os.path.exists(self._be.kernel_modules_dirpath):
            raise self.InstallError("\"%s\" does not exist" % (self._be.kernel_modules_dirpath))
        if not os.path.exists(self._be.firmware_dirpath):
            raise self.InstallError("\"%s\" does not exist" % (self._be.firmware_dirpath))

        # deduplicated disk list
        diskList = OrderedSet()
        for mp in context.get_mount_points():
            for disk in mp.underlay_disk.walk("post-order"):
                diskList.add(disk)

        # get kernel module file list (order is important)
        kmodList = OrderedSet()
        if True:
            kaliasOrderedSet = OrderedSet()

            for disk in diskList:
                if isinstance(disk, HostDiskBtrfsRaid):
                    pass
                elif isinstance(disk, HostDiskBcachefsRaid):
                    pass
                elif isinstance(disk, HostDiskBcache):
                    kaliasOrderedSet.add("bcache")
                elif isinstance(disk, HostDiskLvmLv):
                    kaliasOrderedSet.add("dm_mod")
                elif isinstance(disk, HostDiskScsiHdd):
                    kaliasOrderedSet.add(disk.host_controller_name)
                    kaliasOrderedSet.add("sd_mod")
                elif isinstance(disk, HostDiskNvmeHdd):
                    kaliasOrderedSet.add("nvme")
                elif isinstance(disk, HostDiskXenHdd):
                    kaliasOrderedSet.add("xen_blkfront")
                elif isinstance(disk, HostDiskVirtioHdd):
                    kaliasOrderedSet.add("virtio_pci")
                    kaliasOrderedSet.add("virtio_blk")
                elif isinstance(disk, HostDiskImageFileOnRemovable):
                    # scsi controller drivers, for fixed optical drive / eSATA storage device
                    kaliasOrderedSet.add("ahci")
                    kaliasOrderedSet.add("mpt3sas")
                    # usb stack, for various types of USB storage devices
                    kaliasOrderedSet.add("ehci_pci")
                    kaliasOrderedSet.add("ohci_pci")        # linux kernel warns, ehci_hcd must be before ohci_hcd and uhci_hcd
                    kaliasOrderedSet.add("uhci_hcd")        # and, uhci_hcd soft depends on ehci_pci
                    kaliasOrderedSet.add("xhci_pci")
                    kaliasOrderedSet.add("usb_storage")
                    kaliasOrderedSet.add("uas")
                    # storage type support
                    kaliasOrderedSet.add("sd_mod")          # for SCSI harddisk type of storage
                    kaliasOrderedSet.add("sr_mod")          # for SCSI optical drive
                    kaliasOrderedSet.add("nvme")            # for NVME storage device
                    kaliasOrderedSet.add("mmc_block")       # for SD card type of storage
                    # sg: does it auto get installed, for blueray?
                    # ums-realtek, ums-jmicron: special card reader
                else:
                    assert False

            for disk in diskList:
                if isinstance(disk, HostDiskFeatureUsePartitionTable):
                    if disk.PARTITION_TABLE_TYPE_MBR in disk.used_partition_table_types:
                        # get kernel module for partition format
                        pass
                    if disk.PARTITION_TABLE_TYPE_GPT in disk.used_partition_table_types:
                        # get kernel module for partition format
                        pass

            for disk in diskList:
                if isinstance(disk, HostDiskImageFileOnRemovable):
                    kaliasOrderedSet.add("iso9660")                                                    # cdrom file system
                    kaliasOrderedSet.add("udf")                                                        # cdrom file system
                    self._advancedAddKernelModule(kaliasOrderedSet, "vfat")                            # removable harddisk file system

            for mp in context.get_mount_points():
                if mp.dev_uuid.startswith("UUID:FILE="):
                    kaliasOrderedSet.add("loop")

            for mp in context.get_mount_points():
                if mp.fstype == HostMountPoint.FS_TYPE_VFAT:
                    self._advancedAddKernelModule(kaliasOrderedSet, "vfat")
                elif mp.fstype in [HostMountPoint.FS_TYPE_EXT4, HostMountPoint.FS_TYPE_BTRFS, HostMountPoint.FS_TYPE_BCACHEFS]:
                    kaliasOrderedSet.add(mp.fstype)
                elif mp.fstype == HostMountPoint.FS_TYPE_SQUASHFS_WITH_OVERLAY:
                    kaliasOrderedSet.add("squashfs")
                else:
                    assert False

            for kalias in kaliasOrderedSet:
                kmodList |= OrderedSet(Util.kmodGetKmodAndDeps(self._be.kernel_modules_dirpath, kalias, withDeps=True))

        # get firmware file list
        firmwareList = OrderedSet()
        for km in kmodList:
            firmwareList |= OrderedSet(self._getFirmwareFilepathsByKmod(self._be, km))

        # get block device preparation operation list
        blkOpList = OrderedSet()
        if True:
            bHasBcache = False
            for disk in diskList:
                if isinstance(disk, HostDiskLvmLv):
                    blkOpList.add("lvm-lv-activate %s %s %s" % (disk.uuid, disk.vg_name, disk.lv_name))
                elif isinstance(disk, HostDiskBcache):
                    for cacheDev in disk.cache_dev_list:
                        blkOpList.add("bcache-cache-device-activate %s %s" % (disk.cache_set_uuid, cacheDev.uuid))
                    blkOpList.add("bcache-backing-device-activate %s %s" % (disk.uuid, disk.backing_dev.uuid))
                    bHasBcache = True
                elif isinstance(disk, HostDiskBcachefsRaid):
                    pass
                elif isinstance(disk, HostDiskBtrfsRaid):
                    pass
                elif isinstance(disk, HostDiskScsiHdd):
                    pass
                elif isinstance(disk, HostDiskNvmeHdd):
                    pass
                elif isinstance(disk, HostDiskXenHdd):
                    pass
                elif isinstance(disk, HostDiskVirtioHdd):
                    pass
                elif isinstance(disk, HostDiskImageFileOnRemovable):
                    pass
                else:
                    assert False
            if bHasBcache:
                blkOpList.add("bcache-activate-finish")

        # create basic structure for initramfs directory
        self._installSymlink("usr/bin", "/bin", self._initramfsTmpDir)
        self._installDir("/dev", self._initramfsTmpDir)
        self._installDir("/etc", self._initramfsTmpDir)
        self._installSymlink("usr/lib", "/lib", self._initramfsTmpDir)
        self._installSymlink("usr/lib64", "/lib64", self._initramfsTmpDir)
        self._installDir("/proc", self._initramfsTmpDir)
        self._installDir("/run", self._initramfsTmpDir)
        self._installSymlink("usr/sbin", "/sbin", self._initramfsTmpDir)
        self._installDir("/sys", self._initramfsTmpDir)
        self._installDir("/tmp", self._initramfsTmpDir)
        self._installDir("/usr/bin", self._initramfsTmpDir)
        self._installSymlink("bin", "/usr/sbin", self._initramfsTmpDir)
        self._installDir("/usr/lib", self._initramfsTmpDir)
        self._installDir("/usr/lib64", self._initramfsTmpDir)
        self._installDir("/var", self._initramfsTmpDir)
        self._installDir(self._be.kernel_modules_dirpath, self._initramfsTmpDir)
        self._installDir(self._be.firmware_dirpath, self._initramfsTmpDir)
        self._installDir("/sysroot", self._initramfsTmpDir)
        self._generateDevNode("c", 1, 11, 0x0644, "/dev/kmsg", self._initramfsTmpDir)       # this device is needed before devtmpfs is mounted
        self._generatePasswd(0x0644, "/etc/passwd", self._initramfsTmpDir)
        self._generateGroup(0x0644, "/etc/group", self._initramfsTmpDir)

        # install kmod files
        for f in kmodList:
            Util.copyWithSymlink("/", self._initramfsTmpDir, f[1:])

        # install firmware files
        for f in firmwareList:
            Util.copyWithSymlink("/", self._initramfsTmpDir, f[1:])

        # install files for block device preparation
        # self._installFilesBlkid(self._initramfsTmpDir)
        for disk in diskList:
            if isinstance(disk, HostDiskBtrfsRaid):
                pass
            elif isinstance(disk, HostDiskBcachefsRaid):
                pass
            elif isinstance(disk, HostDiskLvmLv):
                self._installFilesLvm(self._initramfsTmpDir)
            elif isinstance(disk, HostDiskBcache):
                pass
            elif isinstance(disk, HostDiskScsiHdd):
                pass
            elif isinstance(disk, HostDiskNvmeHdd):
                pass
            elif isinstance(disk, HostDiskXenHdd):
                pass
            elif isinstance(disk, HostDiskVirtioHdd):
                pass
            elif isinstance(disk, HostDiskImageFileOnRemovable):
                pass
            else:
                assert False

        # install init executable to initramfs
        self._installBinFromInitDataDir(os.path.join(context.get_tmp_dir(), "init"), self._initramfsTmpDir, "")
        self._installStartupRc(context.get_system_init(), context.get_mount_points(), self._initramfsTmpDir, kmodList, blkOpList)

        # build the initramfs file and tar file
        if True:
            # initramfs file
            cmdStr = "find . -print0 | cpio --null -H newc -o "
            cmdStr += "| xz --format=lzma "            # it seems linux kernel config RD_XZ has bug, so we must use format lzma
            cmdStr += "> \"%s\" " % (self._be.initramfs_filepath)
            subprocess.check_call(cmdStr, shell=True, text=True)

            # tar file
            with tarfile.open(self._be.initramfs_tar_filepath, "w:bz2") as f:
                for fn in glob.glob("*"):
                    f.add(fn)

    def _generateDevNode(self, devType, major, minor, mode, filename, rootDir):
        assert filename.startswith("/")

        dstFile = rootDir + filename
        if devType == "b":
            flag = stat.S_IFBLK
        elif devType == "c":
            flag = stat.S_IFCHR
        else:
            assert False
        os.mknod(dstFile, flag | mode, os.makedev(major, minor))

    def _generatePasswd(self, mode, filename, rootDir):
        assert filename.startswith("/")

        dstFile = rootDir + filename
        with open(dstFile, "w") as f:
            f.write("root:x:0:0::/root:/bin/sh\n")
            f.write("nobody:x:65534:65534::/:/sbin/nologin\n")
        os.chmod(dstFile, mode)

    def _generateGroup(self, mode, filename, rootDir):
        assert filename.startswith("/")

        dstFile = rootDir + filename
        with open(dstFile, "w") as f:
            f.write("tty:x:5:\n")
            f.write("kmem:x:9:\n")
            f.write("disk:x:6:adm\n")
            f.write("floppy:x:11:\n")
            f.write("cdrom:x:19:\n")
        os.chmod(dstFile, mode)

    def _advancedAddKernelModule(self, kaliasOrderedSet, moduleName):
        if moduleName == "vfat":
            buf = pathlib.Path(self._be.kernel_config_filepath).read_text()
            kaliasOrderedSet.add("vfat")
            m = re.search("^CONFIG_FAT_DEFAULT_CODEPAGE=(\\S+)$", buf, re.M)
            if m is not None:
                kaliasOrderedSet.add("nls_cp%s" % (m.group(1)))
            else:
                raise self.InstallError("CONFIG_FAT_DEFAULT_CODEPAGE is missing in \"%s\"" % (self._be.kernel_config_filepath))
            m = re.search("^CONFIG_FAT_DEFAULT_IOCHARSET=\\\"(\\S+)\\\"$", buf, re.M)
            if m is not None:
                kaliasOrderedSet.add("nls_%s" % (m.group(1)))
            else:
                raise self.InstallError("CONFIG_FAT_DEFAULT_IOCHARSET is missing in \"%s\"" % (self._be.kernel_config_filepath))
            return

        assert False

    def _installDir(self, dirPath, rootDir):
        assert dirPath.startswith("/")

        dstDir = rootDir + dirPath
        os.makedirs(dstDir)

    def _installSymlink(self, target, symlinkPath, rootDir):
        assert symlinkPath.startswith("/")

        dstSymlink = rootDir + symlinkPath
        os.symlink(target, dstSymlink)

    def _installBin(self, binFilename, rootDir):
        Util.copyWithSymlink("/", rootDir, binFilename[1:])
        for df in Util.libUsed(binFilename):
            Util.copyWithSymlink("/", rootDir, df[1:])

    def _installBinFromInitDataDir(self, binFilename, rootDir, targetDir):
        dstFilename = os.path.join(rootDir, targetDir, os.path.basename(binFilename))

        subprocess.check_call(["cp", "-f", binFilename, dstFilename])
        subprocess.check_call(["/bin/chmod", "755", dstFilename])

        for df in Util.libUsed(dstFilename):
            Util.copyWithSymlink("/", rootDir, df[1:])

    def _installFilesLvm(self, rootDir):
        self._installBin("/sbin/lvm", self._initramfsTmpDir)

        # note: surrounded " would be recognized as part of rootDir, it's a bug of systemd-tmpfiles
        subprocess.check_call(["/bin/systemd-tmpfiles", "--create", "--root=%s" % (rootDir), "/usr/lib/tmpfiles.d/lvm2.conf"])

        etcDir = os.path.join(rootDir, "etc", "lvm")
        if not os.path.exists(etcDir):
            os.mkdir(etcDir)
        with open(os.path.join(etcDir, "lvm.conf"), "w") as f:
            f.write("global {\n")
            f.write("    locking_type = 4\n")
            f.write("    use_lvmetad = 0\n")
            f.write("}\n")
            f.write("devices {\n")
            f.write("    write_cache_state = 0\n")
            f.write("}\n")
            f.write("backup {\n")
            f.write("    backup = 0\n")
            f.write("    archive = 0\n")
            f.write("}\n")

    def _installStartupRc(self, systemInit, mountPoints, rootDir, kmodList, blkOpList):
        buf = ""

        def _getPrefixedMountPoint(mount_point):
            if mount_point == "/":
                return "./sysroot"
            else:
                assert mount_point.startswith("/")
                return "./sysroot" + mount_point

        # write comments
        for mi in mountPoints:
            buf += "# %s: %s\n" % (mi.mountpoint, mi.dev_uuid)
        buf += "\n"

        # write /run/nologin, it will be deleted by systemd-user-sessions.service
        buf += "echo 'System is booting up. Unprivileged users are not permitted to log in yet. Please come back later. For technical details, see pam_nologin(8).' > /run/nologin\n"
        buf += "\n"

        # load kernel modules
        if len(kmodList) > 0:
            for k in kmodList:
                buf += "insmod \"%s\"\n" % (k)
            buf += "\n"

        # prepare block devices for mount
        if len(blkOpList) > 0:
            for k in blkOpList:
                buf += "%s\n" % (k)
            buf += "\n"

        # prepare directories for mount
        if True:
            bHasLive = False
            for mi in mountPoints:
                if mi.fstype == HostMountPoint.FS_TYPE_SQUASHFS_WITH_OVERLAY and mi.dev_uuid.startswith("UUID:FILE="):
                    assert isinstance(mi.underlay_disk, HostDiskImageFileOnRemovable)
                    assert "UUID:FILE=%s:%s" % (mi.underlay_disk.uuid[len("UUID="):], mi.underlay_disk.image_file_path) == mi.dev_uuid
                    if not bHasLive:
                        bHasLive = True
                        buf += 'mkdir "/run/livedisk"\n'
                        buf += "\n"
                    realUuid = mi.underlay_disk.uuid[len("UUID="):]
                    buf += 'mkdir "/run/livedisk/%s"\n' % (realUuid)
                    buf += 'mkdir "/run/livedisk/%s-squashfs"\n' % (realUuid)
                    buf += 'mkdir "/run/livedisk/%s-overlay"\n' % (realUuid)
                    buf += 'mkdir "/run/livedisk/%s-work"\n' % (realUuid)
                    buf += "\n"

        # FIXME: wait before mount, should wait and retry in /init
        # but after blkid_evaluate_tag is called (first read /dev?), usb stick won't appear in /dev anymore even if it is detected and appeared in /sys, why?
        if True:
            bFound = False
            for mi in mountPoints:
                if mi.dev_uuid.startswith("UUID:FILE="):
                    bFound = True
            if bFound:
                buf += "sleep 10\n"
                buf += "\n"

        # do mount
        for mi in mountPoints:
            if mi.fstype == HostMountPoint.FS_TYPE_SQUASHFS_WITH_OVERLAY and mi.dev_uuid.startswith("UUID:FILE="):
                # don't assert again
                realUuid = mi.underlay_disk.uuid[len("UUID="):]
                buf += 'mount -t vfat %s "/run/livedisk/%s"\n' % (mi.underlay_disk.uuid, realUuid)
                buf += 'mount-file "/run/livedisk/%s-squashfs" squashfs "ro" "/run/livedisk/%s/%s"\n' % (realUuid, realUuid, mi.underlay_disk.image_file_path)
                buf += 'mount-overlayfs "%s" "%s" "/run/livedisk/%s-squashfs" "/run/livedisk/%s-overlay" "/run/livedisk/%s-work"\n' % (_getPrefixedMountPoint(mi.mountpoint), mi.opts, realUuid, realUuid, realUuid)
                buf += "\n"
                continue
            if mi.fstype == "btrfs":
                if isinstance(mi.underlay_disk, HostDiskBtrfsRaid):
                    uuidList = [x.uuid for x in mi.underlay_disk.children]
                else:
                    uuidList = [mi.underlay_disk.uuid]
                buf += 'mount-btrfs "%s" "%s" %s\n' % (_getPrefixedMountPoint(mi.mountpoint), mi.opts, ' '.join(uuidList))
                buf += "\n"
                continue
            if mi.fstype == "bcachefs":
                if isinstance(mi.underlay_disk, HostDiskBcachefsRaid):
                    uuidList = [x.uuid for x in mi.underlay_disk.children]
                else:
                    uuidList = [mi.underlay_disk.uuid]
                buf += 'mount-bcachefs "%s" "%s" %s\n' % (_getPrefixedMountPoint(mi.mountpoint), mi.opts, ' '.join(uuidList))
                buf += "\n"
                continue
            if True:
                buf += 'mount -t %s -o "%s" %s "%s"\n' % (mi.fstype, mi.opts, mi.dev_uuid, _getPrefixedMountPoint(mi.mountpoint))
                buf += "\n"
                continue

        # switch to new root
        buf += ("switchroot \"%s\" %s\n" % (_getPrefixedMountPoint("/"), systemInit.cmd)).rstrip()
        buf += "\n"

        # write cfg file
        with open(os.path.join(rootDir, "startup.rc"), "w") as f:
            f.write(buf)

    @staticmethod
    def _getFirmwareFilepathsByKmod(bootEntry, kmod_filepath):
        ret = []
        for line in subprocess.check_output(["modinfo", kmod_filepath], text=True).split("\n"):
            m = re.fullmatch("firmware: +(\\S.*)", line)
            if m is not None:
                ret.append(os.path.join(bootEntry.firmware_dirpath, m.group(1)))
        return ret
